import sys
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QGraphicsView, QGraphicsScene, QGraphicsProxyWidget, QWidget
from PyQt5.QtCore import Qt, QRectF, QPointF
from PyQt5.QtGui import QImage, QColor, QPainter, QPolygonF

class OccupancyMap:

    def __init__(self, data):
        ### 这里假设data是一个4D numpy数组
        self.data = data  
    
    def GetCostMapData(self):
        return self.data

class DisplayCostMap(QWidget):

    def __init__(self, ):#display_type, z_value, parent_name
        super().__init__()
        self.vehicle_position = (0, 0)  # 初始小车位置
        self.cost_map_data_ = None
        self.map_image_ = None
        self.setGeometry(0, 0, 480, 320) ##710*560

    def update_data(self, data, vehicle_pos=None):
        if isinstance(data, OccupancyMap):
            self.cost_map_data_ = data
            if vehicle_pos is not None:
                self.vehicle_position = vehicle_pos
            self.parse_cost_map()
            self.update()
            return True
        else:
            print("Invalid data type provided.")
            return False

    def parse_cost_map(self):
        cost_map = self.cost_map_data_.GetCostMapData()
        self.map_image_ = QImage(cost_map.shape[1], cost_map.shape[0], QImage.Format_ARGB32)
        
        for i in range(cost_map.shape[1]):
            for j in range(cost_map.shape[0]):
                color_data = cost_map[j, i]
                color = QColor(*color_data)  # 解包RGBA值给QColor
                self.map_image_.setPixelColor(i, j, color)

    def paintEvent(self, event):
        super().paintEvent(event)
        if self.map_image_:
            painter = QPainter(self)
            painter.drawImage(self.rect(), self.map_image_)
            self.draw_vehicle(painter)

    def draw_vehicle(self, painter):
        x, y = self.vehicle_position
        size = 20  # 小车图标大小
        half_size = size // 2
        wheel_radius = size // 4
        head_size = size // 2

        # 绘制车身（矩形）
        body_rect = QRectF(x - half_size, y - size, size, size // 2)
        painter.setBrush(QColor('black'))
        painter.drawRect(body_rect)

        # 绘制左车轮
        left_wheel = QRectF(x - half_size - wheel_radius, y + size // 4 - wheel_radius, 2 * wheel_radius, 2 * wheel_radius)
        painter.setBrush(QColor('black'))
        painter.drawEllipse(left_wheel)

        # 绘制右车轮
        right_wheel = QRectF(x + half_size - wheel_radius, y + size // 4 - wheel_radius, 2 * wheel_radius, 2 * wheel_radius)
        painter.setBrush(QColor('black'))
        painter.drawEllipse(right_wheel)

        # 绘制车头（三角形）
        points = [
            QPointF(x - half_size, y - head_size),
            QPointF(x + half_size, y - head_size),
            QPointF(x, y)
        ]
        painter.setBrush(QColor('red'))
        painter.drawPolygon(QPolygonF(points))
        painter.setBrush(QColor('black'))
        painter.drawEllipse(QRectF(x - half_size, y - head_size, size, size))
        

    def set_vehicle_position(self, pos):
        """设置小车的新位置"""
        self.vehicle_position = pos
        self.update()

class MainWindow(QMainWindow):
    
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Cost Map Display")
        self.setGeometry(100, 100, 800, 600)

        # 创建一个图形视图作为中心部件
        self.view = QGraphicsView(self)
        self.setCentralWidget(self.view)

        # 创建一个场景，并添加DisplayCostMap到场景中
        self.scene = QGraphicsScene(self)
        self.view.setScene(self.scene)

        # 创建DisplayCostMap实例并添加到场景中
        self.display = DisplayCostMap() #"cost_map", 0, "parent"
        proxy_widget = QGraphicsProxyWidget()
        proxy_widget.setWidget(self.display)
        self.scene.addItem(proxy_widget)

        # 模拟数据更新
        self.update_display()

    def update_display(self):
        # 创建一个模拟的成本地图数据
        width, height = 640, 480
        data = np.random.randint(0, 256, (height, width, 4), dtype=np.uint8)
        occupancy_map = OccupancyMap(data)

        # 更新DisplayCostMap的数据
        self.display.update_data(occupancy_map, vehicle_pos=(width//2, height//2))

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())